import logging
import os

from llama_index.core import SimpleDirectoryReader
# 导入llama_index相关模块
from llama_index.core.extractors import SummaryExtractor, QuestionsAnsweredExtractor
from llama_index.core.ingestion import IngestionCache
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.schema import TransformComponent
from llama_index.core.text_splitter import TokenTextSplitter
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)

print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)


# 自定义转换器示例（需要用户实现具体逻辑）
class CustomTransformation(TransformComponent):
    """自定义数据处理组件示例"""

    def __call__(self, nodes, **kwargs):
        # 在此处添加自定义数据处理逻辑
        # 例如：清理数据、添加元数据、转换格式等
        processed_nodes = []
        for node in nodes:
            # 示例：添加自定义元数据
            node.metadata["processing_step2"] = "custom_transform2"
            processed_nodes.append(node)
        return processed_nodes


# 初始化缓存系统（必须在pipeline之前创建）
cache = IngestionCache.from_persist_path("./ingestion_cache.json")

# 创建数据摄入管道
pipeline = IngestionPipeline(
    transformations=[
        # 1. 首先应用自定义转换
        CustomTransformation(),

        # 2. 文本分块处理
        TokenTextSplitter(
            separator=" ",  # 使用空格作为分隔符
            chunk_size=512,  # 每个文本块最大token数
            chunk_overlap=128  # 块之间重叠的token数（保持上下文连贯）
        ),

        # 3. 自动生成摘要
        SummaryExtractor(llm=llm),  # 使用LLM生成每个文本块的摘要

        # 4. 生成问题-答案对
        QuestionsAnsweredExtractor(
            llm=llm,
            questions=3  # 为每个文本块生成3个Q&A对
        )
    ],
    cache=cache  # 启用处理结果缓存
)

# 使用示例（假设有原始文档）
if __name__ == "__main__":
    data_dir = 'D:/pythonRoom/llm_llama/data'
    reader = SimpleDirectoryReader(
        input_dir=data_dir,  # 文件读取
        recursive=True  # 是否递归读取文件中的文件
    )
    documents = reader.load_data()

    # 运行数据处理管道
    processed_nodes = pipeline.run(documents=documents)
    print('==>', processed_nodes)

    # 保存缓存（可选）
    cache.persist("./ingestion_cache.json")
